{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9263c7a2-350f-442b-acce-c4acf568be37",
   "metadata": {},
   "outputs": [],
   "source": [
    "import findspark\n",
    "\n",
    "findspark.init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d764df87-bdf4-4e32-94d4-494dae841dad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on function init in module findspark:\n",
      "\n",
      "init(spark_home=None, python_path=None, edit_rc=False, edit_profile=False)\n",
      "    Make pyspark importable.\n",
      "    \n",
      "    Sets environment variables and adds dependencies to sys.path.\n",
      "    If no Spark location is provided, will try to find an installation.\n",
      "    \n",
      "    Parameters\n",
      "    ----------\n",
      "    spark_home : str, optional, default = None\n",
      "        Path to Spark installation, will try to find automatically\n",
      "        if not provided.\n",
      "    python_path : str, optional, default = None\n",
      "        Path to Python for Spark workers (PYSPARK_PYTHON),\n",
      "        will use the currently running Python if not provided.\n",
      "    edit_rc : bool, optional, default = False\n",
      "        Whether to attempt to persist changes by appending to shell\n",
      "        config.\n",
      "    edit_profile : bool, optional, default = False\n",
      "        Whether to create an IPython startup file to automatically\n",
      "        configure and import pyspark.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(findspark.init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "06e65345-ea4e-46d2-add2-6a24b5a4d527",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "\n",
    "spark = SparkSession.builder.master(\"local[*]\").appName(\"udf_testing\").getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "58b68b5e-9a23-4284-98a8-bf454c76807d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyspark.sql.functions as F\n",
    "import pyspark.sql.types as T\n",
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "06ec99a9-834a-4f1d-811a-7744bd2e6d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "parquet_folder_path = r\"D:\\python work\\learning\\data\\DataAnalysisWithPythonAndPySpark-Data-trunk\\gsod_noaa\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c45e168a-71eb-45c5-9eb2-f38e31c42779",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "folder_list = os.listdir(parquet_folder_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "72fdd583-66ff-4951-8a6c-36b2ce98b477",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_df_from_parquet(file_path):\n",
    "    return spark.read.format(\"parquet\").load(file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "21901308-5cf1-4561-867b-fdfed49a63da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['gsod2010.parquet',\n",
       " 'gsod2011.parquet',\n",
       " 'gsod2012.parquet',\n",
       " 'gsod2013.parquet',\n",
       " 'gsod2014.parquet',\n",
       " 'gsod2015.parquet',\n",
       " 'gsod2016.parquet',\n",
       " 'gsod2017.parquet',\n",
       " 'gsod2018.parquet',\n",
       " 'gsod2019.parquet',\n",
       " 'gsod2020.parquet']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "folder_list "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ad72c0f2-6dc2-4da4-9812-4d9d64ffcd59",
   "metadata": {},
   "outputs": [],
   "source": [
    "gsod = (\n",
    "    reduce(lambda x,y:x.unionByName(x,allowMissingColumns=True),\n",
    "          [read_df_from_parquet(os.path.join(parquet_folder_path,file_path)) for file_path in folder_list[8:10]])\n",
    "    .dropna(subset = [\"year\",\"mo\",\"da\",\"temp\"])\n",
    "    .where(F.col(\"temp\")!=9999.9)\n",
    "    .drop(\"date\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1e8b6c61-11f2-4a3f-b999-46f597f86df8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5ded4453-f3f7-48ca-a49e-50c719b4cd4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@F.pandas_udf(T.DoubleType())\n",
    "def f_to_c(degrees:pd.Series)->pd.Series:\n",
    "    return (degrees-32)*5/9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "359b81ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install PyArrow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e5a8fce7-0b99-4b90-91c5-020113dbc20b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+-------------------+\n",
      "|temp|             temp_c|\n",
      "+----+-------------------+\n",
      "|37.2| 2.8888888888888906|\n",
      "|29.6|-1.3333333333333326|\n",
      "|71.6| 21.999999999999996|\n",
      "|53.5| 11.944444444444445|\n",
      "|24.7| -4.055555555555555|\n",
      "+----+-------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "gsod = gsod.withColumn(\"temp_c\",f_to_c(F.col(\"temp\")))\n",
    "gsod.select(\"temp\",\"temp_c\").distinct().show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e371e44a-f0ac-4759-aa6d-8b0cf7d3fa79",
   "metadata": {},
   "source": [
    "没有管理员权限的电脑，测试不了`F.pandas_udf`会报奇奇怪怪的错误,因为被限制了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8bb297b4-3e8e-4422-bac2-202b17cd56c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+---+---+\n",
      "|year| mo| da|\n",
      "+----+---+---+\n",
      "|2018| 04| 21|\n",
      "|2018| 06| 07|\n",
      "|2018| 03| 12|\n",
      "|2018| 01| 08|\n",
      "|2018| 05| 05|\n",
      "+----+---+---+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "gsod.select(\"year\",\"mo\",\"da\").show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "583046e0-5048-4456-be95-93082d741eee",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.execution.arrow.maxRecodesPerBatch\", 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "243435a4-b7a2-49c4-bc2c-7420b7b6a2dc",
   "metadata": {},
   "source": [
    "设置了批量处理的大小，还是不行"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "4d3b93f9-6cac-4505-8da8-554612332533",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple,Iterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f0dc6ee3-7629-4c84-a450-b49391806d50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+---+---+----------+\n",
      "|year| mo| da|      date|\n",
      "+----+---+---+----------+\n",
      "|2018| 04| 21|2018-04-21|\n",
      "|2018| 06| 07|2018-06-07|\n",
      "|2018| 03| 12|2018-03-12|\n",
      "|2018| 01| 08|2018-01-08|\n",
      "|2018| 05| 05|2018-05-05|\n",
      "+----+---+---+----------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "@F.pandas_udf(T.DateType())\n",
    "def create_date(year_mo_da:Iterator[Tuple[pd.Series,pd.Series,pd.Series]])->Iterator[pd.Series]:\n",
    "    for year,mo,da in year_mo_da:\n",
    "        yield pd.to_datetime(dict(year=year,month=mo,day=da))\n",
    "\n",
    "gsod.select(\"year\",\"mo\",\"da\",create_date(F.col(\"year\"),F.col(\"mo\"),F.col(\"da\")).alias(\"date\")).show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "952c514a",
   "metadata": {},
   "source": [
    "这地方挺奇怪的，参数只设置一个，但是可以传入多个列"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d987aa7c-6f0c-4049-bac7-fccbc4d924ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_date(year_mo_da:Iterator[Tuple[pd.Series,pd.Series,pd.Series]])->Iterator[pd.Series]:\n",
    "    for year,mo,da in year_mo_da:\n",
    "        yield pd.to_datetime(dict(year=year,month=mo,day=da))\n",
    "\n",
    "f = F.pandas_udf(create_date,T.DateType())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "91136ddd-80b5-4eb0-9f36-895c58422c03",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+---+---+----------+\n",
      "|year| mo| da|      date|\n",
      "+----+---+---+----------+\n",
      "|2018| 04| 21|2018-04-21|\n",
      "|2018| 06| 07|2018-06-07|\n",
      "|2018| 03| 12|2018-03-12|\n",
      "|2018| 01| 08|2018-01-08|\n",
      "|2018| 05| 05|2018-05-05|\n",
      "|2018| 01| 27|2018-01-27|\n",
      "|2018| 02| 11|2018-02-11|\n",
      "|2018| 01| 28|2018-01-28|\n",
      "|2018| 01| 28|2018-01-28|\n",
      "|2018| 11| 16|2018-11-16|\n",
      "|2018| 01| 09|2018-01-09|\n",
      "|2018| 04| 07|2018-04-07|\n",
      "|2018| 01| 04|2018-01-04|\n",
      "|2018| 02| 13|2018-02-13|\n",
      "|2018| 01| 13|2018-01-13|\n",
      "|2018| 01| 12|2018-01-12|\n",
      "|2018| 09| 27|2018-09-27|\n",
      "|2018| 10| 12|2018-10-12|\n",
      "|2018| 02| 22|2018-02-22|\n",
      "|2018| 02| 02|2018-02-02|\n",
      "+----+---+---+----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "gsod.select(\"year\",\"mo\",\"da\",f(F.col(\"year\"),F.col(\"mo\"),F.col(\"da\")).alias(\"date\")).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "8adb2ee8-735a-4d90-a2c3-60f17baa2fa9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----+---+-------------------+\n",
      "|   stn|year| mo|        rt_chg_temp|\n",
      "+------+----+---+-------------------+\n",
      "|010010|2018| 06|0.20778642936596214|\n",
      "|010014|2018| 02|-0.1715995568827666|\n",
      "|010014|2018| 05| 0.7449640287769786|\n",
      "|010014|2018| 06|-0.3044848352129638|\n",
      "|010060|2018| 09|-0.4721980771763465|\n",
      "+------+----+---+-------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "@F.pandas_udf(T.DoubleType())\n",
    "def rate_of_change_tempurature(day:pd.Series,temp:pd.Series)->float:\n",
    "    return LinearRegression().fit(X=day.astype(\"int\").values.reshape(-1,1),y=temp).coef_[0]\n",
    "\n",
    "result = gsod.groupby(\"stn\",\"year\",\"mo\").agg(\n",
    "    rate_of_change_tempurature(gsod[\"da\"],gsod[\"temp\"]).alias(\"rt_chg_temp\")\n",
    ")\n",
    "result.distinct().show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b7a6133-959c-4c7c-b3d2-7ad6db05e5fd",
   "metadata": {},
   "source": [
    "普通udf可以做类似apply的操作，不能向pandas_udf那样做分组批量操作"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "6e647e94-e59b-4cba-8a4c-8b3705eafbe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale_temperature(temp_by_day:pd.DataFrame)->pd.DataFrame:\n",
    "    temp = temp_by_day.temp\n",
    "    answer = temp_by_day[[\"stn\",\"year\",\"mo\",\"da\",\"temp\"]]\n",
    "    answer = (answer.assign(temp_norm=0.5) if temp.min() == temp.max() \n",
    "              else answer.assign(temp_norm = (temp-temp.min())/(temp.max()-temp.min())))\n",
    "    return answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "cab9d97b-815e-4650-a2da-f06996be7f95",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----+---+---+----+------------------+\n",
      "|   stn|year| mo| da|temp|         temp_norm|\n",
      "+------+----+---+---+----+------------------+\n",
      "|010014|2018| 02| 04|31.7|0.6818181818181817|\n",
      "|010014|2018| 02| 25|35.6| 0.859090909090909|\n",
      "|010014|2018| 02| 24|36.5|0.8999999999999999|\n",
      "|010014|2018| 02| 18|37.8|0.9590909090909089|\n",
      "|010014|2018| 02| 11|36.8|0.9136363636363634|\n",
      "+------+----+---+---+----+------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "schema = \"stn string, year string, mo string, da string, temp double, temp_norm double\"\n",
    "gsod_map = gsod.groupby(\"stn\",\"year\",\"mo\").applyInPandas(scale_temperature,schema=schema)\n",
    "gsod_map.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d31f8b46",
   "metadata": {},
   "source": [
    "exercise 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "5dffa8eb-be58-4bfa-b827-f62bf41729df",
   "metadata": {},
   "outputs": [],
   "source": [
    "def temp_to_temp(value:pd.Series,from_temp:str,to_temp:str)->pd.Series:\n",
    "    @F.pandas_udf(T.DoubleType())\n",
    "    def c2k(value:pd.Series)->pd.Series:\n",
    "        res_temp = (value*9)/5+32\n",
    "        return res_temp\n",
    "    \n",
    "    if from_temp == \"c\" and to_temp == \"k\":\n",
    "        res = c2k(value)\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "20be353b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----+---+---+----+------------------+------------------+\n",
      "|   stn|year| mo| da|temp|         temp_norm|            temp_K|\n",
      "+------+----+---+---+----+------------------+------------------+\n",
      "|010014|2018| 02| 04|31.7|0.6818181818181817|             89.06|\n",
      "|010014|2018| 02| 25|35.6| 0.859090909090909| 96.08000000000001|\n",
      "|010014|2018| 02| 24|36.5|0.8999999999999999|              97.7|\n",
      "|010014|2018| 02| 18|37.8|0.9590909090909089|100.03999999999999|\n",
      "|010014|2018| 02| 11|36.8|0.9136363636363634|             98.24|\n",
      "+------+----+---+---+----+------------------+------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "gsod_map.withColumn(\"temp_K\",temp_to_temp(F.col(\"temp\"),\"c\",\"k\")).show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1473018",
   "metadata": {},
   "source": [
    "exercise 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "1537ab14",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----+---+---+----+------------------+\n",
      "|   stn|year| mo| da|temp|            temp_K|\n",
      "+------+----+---+---+----+------------------+\n",
      "|010014|2018| 02| 04|31.7|             89.06|\n",
      "|010014|2018| 02| 25|35.6| 96.08000000000001|\n",
      "|010014|2018| 02| 24|36.5|              97.7|\n",
      "|010014|2018| 02| 18|37.8|100.03999999999999|\n",
      "|010014|2018| 02| 11|36.8|             98.24|\n",
      "+------+----+---+---+----+------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def scale_temperature(temp_by_day:pd.DataFrame)->pd.DataFrame:\n",
    "    temp = temp_by_day.temp\n",
    "    answer = temp_by_day[[\"stn\",\"year\",\"mo\",\"da\",\"temp\"]]\n",
    "    answer = answer.assign(temp_K=(temp*9)/5+32)\n",
    "    return answer\n",
    "\n",
    "schema = \"stn string, year string, mo string, da string, temp double, temp_K double\"\n",
    "gsod.groupby(\"stn\",\"year\",\"mo\").applyInPandas(scale_temperature,schema=schema).show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02348de8",
   "metadata": {},
   "source": [
    "exercise 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "d545f675",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "2d83f53a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----+---+--------------------+\n",
      "|   stn|year| mo|         rt_chg_temp|\n",
      "+------+----+---+--------------------+\n",
      "|010010|2018| 06|[0.20778642936596...|\n",
      "|010014|2018| 02|[-0.1715995568827...|\n",
      "|010014|2018| 05|[0.74496402877697...|\n",
      "|010014|2018| 06|[-0.3044848352129...|\n",
      "|010060|2018| 09|[-0.4721980771763...|\n",
      "+------+----+---+--------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "@F.pandas_udf(T.ArrayType(T.DoubleType()))\n",
    "def rate_of_change_tempurature(day:pd.Series,temp:pd.Series)->List[float]:\n",
    "    lr =  LinearRegression()\n",
    "    lr.fit(X=day.astype(\"int\").values.reshape(-1,1),y=temp)\n",
    "    return [lr.coef_[0],lr.intercept_]\n",
    "\n",
    "result = gsod.groupby(\"stn\",\"year\",\"mo\").agg(\n",
    "    rate_of_change_tempurature(gsod[\"da\"],gsod[\"temp\"]).alias(\"rt_chg_temp\")\n",
    ")\n",
    "result.distinct().show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "d50cfac7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----+---+-----------------------------------------+\n",
      "|stn   |year|mo |rt_chg_temp                              |\n",
      "+------+----+---+-----------------------------------------+\n",
      "|010010|2018|06 |[0.20778642936596214, 34.94597701149424] |\n",
      "|010014|2018|02 |[-0.1715995568827666, 36.100857335516814]|\n",
      "|010014|2018|05 |[0.7449640287769786, 48.64234326824254]  |\n",
      "|010014|2018|06 |[-0.3044848352129638, 62.43070842001532] |\n",
      "|010060|2018|09 |[-0.4721980771763465, 41.450428025813245]|\n",
      "+------+----+---+-----------------------------------------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "result.distinct().show(5,False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d171ee3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
